import torch
# from torch.utils.data import DataLoader, TensorDataset
# import pandas as pd
import numpy as np
from data_process import *
from train_demo2 import *
from utility.import_log import flog
log = flog(__file__)

# 加载模型
def load_model():
    log.info('开始加载模型文件')
    try:
        model=torch.load('model.pkl')
        return model
    except Exception as e:
        log.error(f'模型加载失败:{e}')



# 预测
def predict(seq):
    index = np.expand_dims(seq2index(seq), 0)
    x = padding_seq(index)
    model = load_model()
    pred_output = model(torch.from_numpy(x))
    pred_y = torch.argmax(pred_output[0]).data.numpy()
    return pred_y

    # df_chunksize = pd.read_csv('predict_data.csv', header=None, chunksize=batch_size)
    # chunk_num=0
    # for df in df_chunksize:
    #     df_x = df[0]
    #     df_x = padding_seq(df_x.apply(seq2index))
    #     model = load_model()
    #     pred_output = model(torch.from_numpy(df_x))
    #     pred_y = torch.max(pred_output, 1)[1].data.numpy()
    #     pred = pd.DataFrame({'question':df[0].values,'pred':pred_y})
    #     pred.to_csv('predict_result.csv', mode='a',header=None,index=None)
    #     chunk_num = chunk_num+1
    #     print('chunk:',chunk_num,'  预测已完成')



if __name__ == '__main__':
    y=predict('薯条的身份证号是什么？')
    print(y)

